import jsonlines
import json
import random


def make_conversation_list(file_path):
    f = list(jsonlines.open(file_path,'r'))
    result = []
    for index,line in enumerate(f):
        data = []
        flag = False
        conversations = line['conversation'].split('\n\n')
        for conversation in conversations:
            if "老师：" in conversation:
                lines = conversation.split('老师：')
                role = "老师"
            elif "学生：" in conversation:
                lines = conversation.split('学生：')
                role = "学生"
            else:
                print(conversation)
                break
                flag=True
            message = lines[1]
            print(role)
            print(message)
            data.append({
                "from":role,
                "value":message
            })
        if flag:
            continue
        else:
            result.append({
                "id":f"identity_{index}",
                "conversations": data
            })
    return result


def make_conversation_list2(file_path):
    f= list(jsonlines.open(file_path,'r'))
    result = []
    for index,line in enumerate(f):
        data = []
        flag = False
        conversation = "\n\n"+line['conversation']
        conversation = conversation.split('\n\n')[1:]
        for c in conversation:
            role = c[:3][:2]
            sentence = c[3:]
            if role not in ("老师","学生"):
                flag = True

            if "：" in sentence:
                flag = True

            data.append({
                "from":role,
                "value":sentence
            })
        if flag:
            continue
        else:
            result.append({
                "id":f"identity_{index}",
                "conversations": data
            })
    
    return result
        
            



if __name__ == "__main__":
    file_path0 = '/mnt/lth/llm/make_dataset/q3_2.jsonl'
    data_list0 = make_conversation_list2(file_path0)
    file_path1 = "/mnt/lth/llm/make_dataset/q3_3.jsonl"
    data_list1 = make_conversation_list2(file_path1)

    data_list = data_list0+data_list1
    random.shuffle(data_list)
    f_path = 'q3-convert.json'
    with open(f_path, 'w', encoding='utf-8') as f:
        json.dump(data_list, f, ensure_ascii=False, indent=4)


    # file_path =  '/mnt/lth/llm/make_dataset/q3_2.jsonl'
    # data_list = check_conversation(file_path)